/*
 * Copyright (c) 2020-2021. the original authors and DEPSEA.ORG
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */
package org.depsea.log.advice;

import brave.Span;
import brave.Tracer;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.module.SimpleModule;
import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.exception.ExceptionUtils;
import org.depsea.log.exception.LogPointAspectException;
import org.depsea.log.getter.ApiNameGetter;
import org.depsea.log.getter.UserGetter;
import org.depsea.log.handler.RequestPointHandler;
import org.depsea.log.model.RequestInfo;
import org.depsea.log.notification.ExceptionNotification;
import org.depsea.log.serializer.MultipartFileSerializer;
import org.depsea.log.utils.IpUtils;
import org.springframework.aop.AfterReturningAdvice;
import org.springframework.aop.MethodBeforeAdvice;
import org.springframework.aop.ThrowsAdvice;
import org.springframework.beans.BeansException;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.lang.NonNull;
import org.springframework.util.CollectionUtils;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;
import org.springframework.web.multipart.MultipartFile;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.lang.reflect.Method;
import java.lang.reflect.Parameter;
import java.util.*;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.stream.Collectors;

/**
 * @author jaune
 * @since 1.0.0
 */
@Slf4j
public class RequestAdvice implements AfterReturningAdvice, ThrowsAdvice, MethodBeforeAdvice, ApplicationContextAware {

    /**
     * 线程池，用于执行异步处理
     */
    private final ExecutorService executorService = Executors.newCachedThreadPool();

    /**
     * 是否包含请求头信息
     */
    private boolean includeRequestHeaders = true;

    /**
     * 是否包含应答头信息
     */
    private boolean includeResponseHeaders = true;

    private ApiNameGetter apiNameGetter;

    private UserGetter userGetter;

    /**
     * 需要 spring-cloud-starter-sleuth 支持，用于获取链路追踪信息
     */
    private Tracer tracer;

    /**
     * 用于记录耗时
     */
    private final ThreadLocal<Long> timeConsumingThreadLocal = new ThreadLocal<>();

    /**
     * JSON序列化
     */
    private final ObjectMapper objectMapper = new ObjectMapper();

    /**
     * 请求信息处理
     */
    private List<RequestPointHandler> requestPointHandlers = new ArrayList<>();

    /**
     * 异常请求通知
     */
    private List<ExceptionNotification> exceptionNotifications = new ArrayList<>();

    /**
     * Spring 上下文
     */
    private ApplicationContext applicationContext;

    private final String localIp;

    public void setApiNameGetter(ApiNameGetter apiNameGetter) {
        this.apiNameGetter = apiNameGetter;
    }

    public void setTracer(Tracer tracer) {
        this.tracer = tracer;
    }

    public void setIncludeRequestHeaders(boolean includeRequestHeaders) {
        this.includeRequestHeaders = includeRequestHeaders;
    }

    public void setIncludeResponseHeaders(boolean includeResponseHeaders) {
        this.includeResponseHeaders = includeResponseHeaders;
    }

    public void setPointHandlers(List<RequestPointHandler> requestPointHandlers) {
        this.requestPointHandlers = requestPointHandlers;
    }

    public void addPointHandlers(RequestPointHandler... requestPointHandlers) {
        this.requestPointHandlers.addAll(Arrays.stream(requestPointHandlers).collect(Collectors.toList()));
    }

    public void setExceptionNotifications(List<ExceptionNotification> exceptionNotifications) {
        this.exceptionNotifications = exceptionNotifications;
    }

    public void addExceptionNotification(ExceptionNotification... exceptionNotifications) {
        this.exceptionNotifications.addAll(Arrays.stream(exceptionNotifications).collect(Collectors.toList()));
    }

    public void setUserGetter(UserGetter userGetter) {
        this.userGetter = userGetter;
    }

    public RequestAdvice() {
        Set<String> localIpSet = IpUtils.getLocalIp();
        this.localIp = StringUtils.join(localIpSet, ",");
        this.objectMapper.registerModule(new JavaTimeModule());

        SimpleModule module = new SimpleModule();
        module.addSerializer(MultipartFile.class, new MultipartFileSerializer());
        this.objectMapper.registerModule(module);
    }

    @Override
    public void setApplicationContext(@NonNull ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
    }

    @Override
    public void before(@NonNull Method method, @NonNull Object[] args, Object target) throws Throwable {
        // record start time
        this.timeConsumingThreadLocal.remove();
        this.timeConsumingThreadLocal.set(System.currentTimeMillis());
    }

    @Override
    public void afterReturning(Object returnValue, @NonNull Method method, @NonNull Object[] args, Object target) throws Throwable {
        try {
            RequestInfo requestPoint = this.createRequestPoint(method, args);
            if (returnValue != null) {
                requestPoint.setReturnValue(this.objectMapper.writeValueAsString(returnValue));
            }
            this.doRequestHandler(requestPoint);
        } catch (Exception ex) {
            log.warn("处理失败：{}", ex.getMessage(), ex);
        }
    }

    public void afterThrowing(Method method, Object[] args, Object target, Exception ex) {
        try {
            RequestInfo requestPoint = this.createRequestPoint(method, args);
            requestPoint.setError(true);
            requestPoint.setExceptionName(ex.getClass().getName());
            requestPoint.setExceptionStack(ExceptionUtils.getStackTrace(ex));
            requestPoint.setErrorMessage(ex.getMessage());
            this.doRequestHandler(requestPoint);
            this.doExceptionNotificationHandler(ex, requestPoint);
        } catch (Exception e) {
            log.warn("处理失败：{}", e.getMessage(), ex);
        }
    }

    private RequestInfo createRequestPoint(@NonNull Method method, @NonNull Object[] args) throws JsonProcessingException {
        HttpServletRequest request = this.getRequest();
        HttpServletResponse response = this.getResponse();

        RequestInfo requestInfo = new RequestInfo();

        if (this.includeRequestHeaders) {
            requestInfo.setRequestHeaders(this.objectMapper.writeValueAsString(this.getRequestHeaders(request)));
        }
        if (this.includeResponseHeaders && response != null) {
            requestInfo.setResponseHeaders(this.objectMapper.writeValueAsString(this.getResponseHeaders(response)));
        }

        if (this.apiNameGetter != null) {
            requestInfo.setModuleName(this.apiNameGetter.getModuleName(method.getDeclaringClass()));
            requestInfo.setApiName(this.apiNameGetter.getOperationName(method));
        }

        if (this.userGetter != null) {
            requestInfo.setUserId(this.userGetter.getUserId());
            requestInfo.setName(this.userGetter.getName());
        }

        requestInfo.setServerLocalIp(this.localIp);
        requestInfo.setRequestIpAddress(IpUtils.getRequestIp(request));
        requestInfo.setApplicationName(this.getApplicationName());
        requestInfo.setClazz(method.getDeclaringClass().getName());
        requestInfo.setMethodName(method.getName());
        requestInfo.setError(false);
        requestInfo.setTimestamp(new Date());

        if (this.tracer != null) {
            Span span = tracer.currentSpan();
            if (span != null) {
                requestInfo.setTraceId(Long.toHexString(span.context().traceId()));
                requestInfo.setTraceIdNum(span.context().traceId());
                requestInfo.setSpanId(Long.toHexString(span.context().spanId()));
                requestInfo.setSpanIdNum(span.context().spanId());
                if (span.context().parentId() != null) {
                    requestInfo.setParentId(Long.toHexString(span.context().parentId()));
                }
            }
        }

        requestInfo.setSchema(request.getScheme().toUpperCase());
        requestInfo.setRequestMethod(request.getMethod().toUpperCase());
        requestInfo.setRequestUri(request.getRequestURI());

        Map<String, Object> methodParams = this.getParameterMap(method, args);

        requestInfo.setMethodParameterMap(objectMapper.writeValueAsString(methodParams));
        requestInfo.setRequestParameterMap(objectMapper.writeValueAsString(request.getParameterMap()));
        if (response != null) {
            requestInfo.setResponseStatus(response.getStatus());
        }

        long start = this.timeConsumingThreadLocal.get();
        long end = System.currentTimeMillis();
        requestInfo.setTimeConsuming(end - start);
        this.timeConsumingThreadLocal.remove();
        return requestInfo;
    }

    private String getApplicationName() {
        return this.applicationContext.getEnvironment().getProperty("spring.application.name");
    }

    private Map<String, Object> getParameterMap(Method method, Object[] args) {
        Map<String, Object> methodParams = new HashMap<>();
        Parameter[] parameters = method.getParameters();
        for (int i = 0; i < parameters.length; i++) {
            if (args[i] instanceof HttpServletRequest || args[i] instanceof HttpServletResponse) {
                continue;
            }
            Parameter parameter = parameters[i];
            methodParams.put(parameter.getName(), args[i]);
        }
        return methodParams;
    }

    private Map<String, String> getRequestHeaders(HttpServletRequest request) {
        Map<String, String> headers = new HashMap<>();
        Enumeration<String> headerNames = request.getHeaderNames();
        while (headerNames.hasMoreElements()) {
            String headerName = headerNames.nextElement();
            headers.put(headerName, request.getHeader(headerName));
        }
        return headers;
    }

    private Map<String, String> getResponseHeaders(HttpServletResponse response) {
        Map<String, String> headers = new HashMap<>();
        Collection<String> headerNames = response.getHeaderNames();
        if (CollectionUtils.isEmpty(headerNames)) {
            return headers;
        }
        for (String headerName : headerNames) {
            headers.put(headerName, response.getHeader(headerName));
        }
        return headers;
    }

    private HttpServletRequest getRequest() {
        Optional<RequestAttributes> requestAttributesOptional = Optional.ofNullable(RequestContextHolder.getRequestAttributes());
        if (requestAttributesOptional.isPresent()) {
            ServletRequestAttributes servletRequestAttributes = (ServletRequestAttributes) requestAttributesOptional.get();
            return servletRequestAttributes.getRequest();
        } else {
            throw new LogPointAspectException("Could not get the HttpServletRequest from the spring webmvc context.");
        }
    }

    private HttpServletResponse getResponse() {
        Optional<RequestAttributes> requestAttributesOptional = Optional.ofNullable(RequestContextHolder.getRequestAttributes());
        if (requestAttributesOptional.isPresent()) {
            ServletRequestAttributes servletRequestAttributes = (ServletRequestAttributes) requestAttributesOptional.get();
            return servletRequestAttributes.getResponse();
        } else {
            throw new LogPointAspectException("Could not get the HttpServletResponse from the spring webmvc context.");
        }
    }

    private void doRequestHandler(RequestInfo requestInfo) {
        if (!CollectionUtils.isEmpty(this.requestPointHandlers)) {
            for (RequestPointHandler requestPointHandler : this.requestPointHandlers) {
                try {
                    if (requestPointHandler.isAsync()) {
                        this.executorService.submit(() -> requestPointHandler.handle(requestInfo));
                    } else {
                        requestPointHandler.handle(requestInfo);
                    }
                } catch (Exception ex) {
                    log.warn("Execute request point handler [{}] fail. Cause: {}",
                            requestPointHandler.getClass().getName(), ex.getMessage(), ex);
                }
            }
        }
    }

    private void doExceptionNotificationHandler(Exception ex, RequestInfo requestInfo) {
        if (!CollectionUtils.isEmpty(this.exceptionNotifications)) {
            for (ExceptionNotification exceptionNotification : exceptionNotifications) {
                this.executorService.submit(() -> exceptionNotification.doNotify(ex, requestInfo));
            }
        }
    }

}